package com.shujia.kafka;


import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.java.tuple.Tuple;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.GlobalWindows;
import org.apache.flink.streaming.api.windowing.evictors.CountEvictor;
import org.apache.flink.streaming.api.windowing.triggers.CountTrigger;
import org.apache.flink.streaming.api.windowing.windows.GlobalWindow;
import org.apache.flink.util.Collector;

import java.util.Comparator;
import java.util.PriorityQueue;

import static org.apache.flink.streaming.api.environment.CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION;

public class test {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.enableCheckpointing(5000);
        env.getCheckpointConfig().setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);

        // 使用 externalized checkpoints，这样 checkpoint 在作业取消后仍就会被保留
        env.getCheckpointConfig().setExternalizedCheckpointCleanup(RETAIN_ON_CANCELLATION);

        env.setStateBackend(new HashMapStateBackend());
        KafkaSource<String> src = KafkaSource.<String>builder()
                .setBootstrapServers("hadoop102:9092,hadoop103:9092,hadoop104:9092")
                .setGroupId("StudentDataReader")
                .setTopics("student")
                .setValueOnlyDeserializer(new SimpleStringSchema())
                .build();
        DataStreamSource<String> kafkaDS = env.fromSource(src, WatermarkStrategy.noWatermarks(), "student");

        DataStream<Tuple2<String, Integer>> wordCounts = kafkaDS
                .flatMap(new Tokenizer())
                .keyBy(kv -> kv.f0)
                .sum(1);
        wordCounts.print();
//
        SingleOutputStreamOperator<Tuple2<String, Integer>> sortedWordCounts = wordCounts
                .keyBy(kv -> kv.f0)
                .window(GlobalWindows.create())
                .trigger(CountTrigger.of(1))
                .evictor(CountEvictor.of(0))
                .aggregate(new TopNFunction(), new ProcessWindowFunction<Tuple2<String, Integer>, Tuple2<String, Integer>, String, GlobalWindow>() {
                    @Override
                    public void process(String s, ProcessWindowFunction<Tuple2<String, Integer>, Tuple2<String, Integer>, String, GlobalWindow>.Context context, Iterable<Tuple2<String, Integer>> iterable, Collector<Tuple2<String, Integer>> collector) throws Exception {
                        for (Tuple2<String, Integer> tuple2 : iterable) {
                            collector.collect(tuple2);
                        }
                    }
                });

        sortedWordCounts.print();

//
//
//
        env.execute();
//    }
//
    }

    public static class Tokenizer implements FlatMapFunction<String, Tuple2<String, Integer>> {
        @Override
        public void flatMap(String value, Collector<Tuple2<String, Integer>> out) {
            // normalize and split the line into words
            String[] words = value.toLowerCase().split(",");

            // emit the words
            for (String word : words) {
                out.collect(new Tuple2<>(word, 1));
            }
        }
    }

    public static class TopNFunction implements AggregateFunction<Tuple2<String, Integer>, PriorityQueue<Tuple2<String, Integer>>, Tuple2<String, Integer>> {
        @Override
        public PriorityQueue<Tuple2<String, Integer>> createAccumulator() {
            return new PriorityQueue<>(Comparator.comparingInt(kv-> kv.f1));
        }

        @Override
        public PriorityQueue<Tuple2<String, Integer>> add(Tuple2<String, Integer> value, PriorityQueue<Tuple2<String, Integer>> accumulator) {
            accumulator.offer(value);
            if (accumulator.size() > 100) {
                accumulator.poll();
            }
            return accumulator;
        }

        @Override
        public Tuple2<String, Integer> getResult(PriorityQueue<Tuple2<String, Integer>> accumulator) {
            Tuple2<String, Integer> result = accumulator.poll();
            return result;
        }

        @Override
        public PriorityQueue<Tuple2<String, Integer>> merge(PriorityQueue<Tuple2<String, Integer>> a, PriorityQueue<Tuple2<String, Integer>> b) {
            PriorityQueue<Tuple2<String, Integer>> merged = new PriorityQueue<>(Comparator.comparingInt(kv-> kv.f1));
            merged.addAll(a);
            merged.addAll(b);
            while (merged.size() > 100) {
                merged.poll();
            }
            return merged;
        }
    }

}
